#!/usr/bin/env python2.7

import os, sys
import re
import numpy as np
from StringIO import StringIO
from metlib.datetime import *
from metlib import grep

__all__ = ['read_hkwpf_hourly', 'read_hkwpf', 'hkwpf_getarr']

hkwpf_hourly_dtype = np.dtype([
    ('qc', 'i4'), 
    ('height', 'f8'),
    ('ws', 'f8'),
    ('wd', 'f8'),
    ('u', 'f8'),
    ('v', 'f8'),
    ('w', 'f8'),
    ('noNE', 'i4'),
    ('noSE', 'i4'),
    ('noV', 'i4'),
    ('snrNE', 'f8'),
    ('snrSE', 'f8'),
    ('snrV', 'f8')
    ])

hkwpf_hourly_pattern = r'^ \d\s+\d+\s+[-\.\d]+\s+[-\.\d]+\s+[-\.\d]+\s+[-\.\d]+\s+[-\.\d]+\s+\d+\s+\d+\s+\d+\s+[-\.\d]+\s+[-\.\d]+\s+[-\.\d]+$'

def read_hkwpf_hourly(fname):
    bname = os.path.basename(fname)
    stname = bname[0:3]
    hl = bname[3:5]
    dt = T(bname[5:])
    
    rawlines = [l.rstrip() for l in open(fname).readlines()]
    validlines = grep(hkwpf_hourly_pattern, rawlines)
    line = '\n'.join(validlines)
    res = np.loadtxt(StringIO(line), dtype=hkwpf_hourly_dtype)
    for field, tp in hkwpf_hourly_dtype.descr:
        if tp in [ '<f8', '<f4']:
            res[field][res[field]<= -950.] = np.nan
    if res[0]['height'] > 1000.0:
        return res[1:]
    else:
        return res

def read_hkwpf(rootdir, stname, beg, end, delta='1h'):
    dts = datetime_range(beg, end, delta)
    fnames = ['%s/%s/%s/%s%s' % (rootdir, dt.strftime('%Y'), dt.strftime('%Y%m'), stname, dt.strftime('%Y%m%d%H%M')) for dt in dts]
    reslist = []
    for fname in fnames:
        print fname
        try:
            prof = read_hkwpf_hourly(fname)
        except Exception as e:
            prof = None
        reslist.append(prof)
    return dts, reslist
    
def hkwpf_getarr(reslist, varname):
    reclens = []
    resnum = len(reslist)
    
    for rec in reslist:
        if rec is not None:
            reclens.append(len(rec))
        else:
            reclens.append(0)
    wvalid = np.where(reclens != 0)


    reclen = reclens[1]
    res = np.zeros((resnum, reclen))
    res[:] = np.nan
    for i,rec in enumerate(reslist):
        if rec is not None:
            res[i][:] = rec[varname]

    return res








